import os

import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import minimize

FctrSymtry = 2.0


# ========================================= 将 Dose 归一化, 并转换成 密度分布
def normalize_by_sum(arr, binwidth):
    sum = np.sum(arr)
    if sum == 0:
        return np.zeros_like(arr)
    return arr / sum / binwidth


def weighted_stddev(x, y, eps=1e-12):
    y = np.maximum(y, eps)
    w = y / np.sum(y)
    mu = np.sum(w * x)
    var = np.sum(w * (x - mu) ** 2)
    return np.sqrt(var)


def multi_gauss_model(x, params, n_gauss, with_background=False):
    """
    参数排列: [normFactor, w2..wn, s1..sn, (c)]
    normFactor: 总幅度 scale factor
    wi: 各高斯分量权重; 总和为1,所以省略第一个参数
    si: 各高斯分量sigma;
    c : 常数背景项 (可选)
    """
    normFactor = params[0]
    w = params[1:n_gauss]
    s = params[n_gauss : 2 * n_gauss]
    y_pred = np.zeros_like(x, dtype=float)

    w1 = 1 - np.sum(w)
    y_pred += FctrSymtry * w1 * np.exp(-0.5 * (x / s[0]) ** 2) / (np.sqrt(2 * np.pi) * s[0])
    for wi, si in zip(w, s[1:]):
        y_pred += FctrSymtry * wi * np.exp(-0.5 * (x / si) ** 2) / (np.sqrt(2 * np.pi) * si)
    y_pred *= normFactor
    if with_background:
        y_pred += params[-1]
    return y_pred


def sigma_smooth_penalty(s, s_prev, lam_s=1e3):
    return lam_s * np.sum(((s - s_prev) / np.maximum(s_prev, 1e-3)) ** 2)


def mid_tail_weights(x, center=15.0, width=3.0, lam_m=10.0):
    """
    在 |x|位于center附近时放大权重, 来改善这个区域的拟合. 指数权重势
    """
    r = np.abs(x)
    return 1.0 + lam_m * np.exp(-0.5 * ((x - center) / max(width, 1e-6)) ** 2)


## x, y = Yaxis, profile
def mid_tail_detect(
    x,
    y,
    y_pred0,
    r_min=5.0,
    r_max=30.0,
    smooth_window=10,
    width_clip=(3.0, 12.0),
    lam_m_clip=(6.0, 30.0),
):
    """
    初步拟合得到 y_pred0, 然后根据误差情况自动配置 (center, width, lam_m).
    - 在 [r_min, r_max] 的范围内寻找误差最大的位置;
    - width 为 HWHM, 并裁剪到 width_clip;
    - lam_m 与最大相对误差线性映射到 lam_m_clip.
    """
    eps = 1e-12
    r = np.abs(x)
    rltv_err = np.abs(y_pred0 - y) / np.maximum(y, eps)

    ## 只取主峰两侧的中间部分
    mask = (r >= r_min) & (r <= r_max)
    if not np.any(mask):
        return dict(center=15.0, width=5.0, lam_m=1.0)

    r_mask = r[mask]
    rltv_err_mask = rltv_err[mask]

    ## 因为涉及到正负两侧的range, sort并做简易平滑; _s for sorted
    idx = np.argsort(r_mask)
    r_s = r_mask[idx]
    rltv_err_s = rltv_err_mask[idx]
    ## 平滑, _sm for smoothed
    if smooth_window > 1 and smooth_window < rltv_err_s.size:
        k = smooth_window
        kernel = np.ones(k, dtype=float) / k
        rltv_sm = np.convolve(rltv_err_s, kernel, mode="same")
    else:
        rltv_sm = rltv_err_s

    ## 误差最大位置
    idx_peak = int(np.argmax(rltv_sm))
    center = r_s[idx_peak]

    ## width
    half = 0.5 * rltv_sm[idx_peak]
    ## 左侧
    iL = 0
    for i in range(idx_peak, -1, -1):
        if rltv_sm[i] < half:
            iL = i
            break
    ## 右侧
    iR = len(rltv_sm) - 1
    for i in range(idx_peak, len(rltv_sm)):
        if rltv_sm[i] < half:
            iR = i
            break
    width = 0.5 * max(1e-6, r_s[iR] - r_s[iL])
    width = np.clip(width, width_clip[0], width_clip[1])

    max_rel = rltv_sm[idx_peak]
    print(f"max_rel = {max_rel:.4f}")
    lam_m = 8.0 + 18.0 * (max_rel / (max_rel + 0.1))
    lam_m = np.clip(lam_m, lam_m_clip[0], lam_m_clip[1])

    return dict(center=center, width=width, lam_m=lam_m)


def mid_tail_detect_v2(
    x,
    y,
    y_pred0,
    r_min=5.0,
    r_max=40.0,
    smooth_window=15,
    width_clip=(2.0, 40.0),
    lam_m_clip=(6.0, 40.0),
    rel_err=0.01,
    abs_err=0.0,
):
    """
    查找主峰两侧的中尾部 y_pred < y 的区域.
    """
    eps = 1e-12
    diff = y_pred0 - y
    thres = np.maximum(abs_err, rel_err * np.maximum(y, eps))

    ## y_pred < y 部分
    sel_mask = (diff < -thres) & (np.abs(x) >= r_min) & (np.abs(x) <= r_max)
    if not np.any(sel_mask):
        return dict(x_left=None, x_right=None, center=15.0, width=5.0, lam_m=1.0)

    x_slctd = x[sel_mask]
    rltv_slctd = -(diff[sel_mask]) / np.maximum(y[sel_mask], eps)
    ## sorted
    idx = np.argsort(x_slctd)
    x_std = x_slctd[idx]
    rel_std = rltv_slctd[idx]

    ## 找出最长的连续 y_pred < y 部分
    if len(x_std) < 3:
        return dict(
            x_left=x_std.min(),
            x_right=x_std.max(),
            center=np.mean(x_std),
            width=2.0,
            lam_m=8.0,
        )

    dx = 0.1  # 相邻点横坐标间距
    gap_thres = 1.5 * dx

    ## 如果相邻点间距大于gap_thres则认为不连续;
    ## 连续区间的 start, end pairs;
    starts = [0]
    for i in range(1, len(x_std)):
        if (x_std[i] - x_std[i - 1]) > gap_thres:
            starts.append(i)
    ends = starts[1:] + [len(x_std)]
    ## search 最长的连续区间
    seg_lengths = [ends[k] - starts[k] for k in range(len(starts))]
    k_max = int(np.argmax(seg_lengths))
    left, right = starts[k_max], ends[k_max]
    x_seg, rel_seg = x_std[left:right], rel_std[left:right]

    if len(x_seg) < 3:
        return dict(
            x_left=x_seg.min(),
            x_right=x_seg.max(),
            center=np.mean(x_seg),
            width=2.0,
            lam_m=8.0,
        )

    # 简单平滑并确定误差最大位置
    if smooth_window > 1 and smooth_window < len(rel_seg):
        kernel = np.ones(smooth_window) / smooth_window
        rel_sm = np.convolve(rel_seg, kernel, mode="same")
    else:
        rel_sm = rel_seg

    # 找y_pred < y 部分的center的位置
    i_peak = np.argmax(rel_sm)
    x_peak = x_seg[i_peak]
    peak_val = rel_sm[i_peak]
    print(f"peak_val: {peak_val:.3f}")

    # 定位width
    half = 0.1 * peak_val
    iL, iR = 0, len(x_seg) - 1
    for i in range(i_peak, -1, -1):
        if rel_sm[i] < half:
            iL = i
            break
    for i in range(i_peak, len(x_seg)):
        if rel_sm[i] < half:
            iR = i
            break

    x_left, x_right = x_seg[iL], x_seg[iR]
    width = 0.5 * abs(x_right - x_left)
    width = np.clip(width, width_clip[0], width_clip[1])

    ## 权重强度
    # lam_m = 8.0 + 18.0 * (peak_val / (peak_val + 0.2))
    lam_m = 8.0 + 30.0 * (1 - np.exp(-3.0 * peak_val))
    lam_m = np.clip(lam_m, lam_m_clip[0], lam_m_clip[1])

    return dict(center=x_peak, width=width, lam_m=lam_m)


def objective_log(
    params,
    x,
    y,
    n_gauss,
    with_background,
    lam_n=1e3,
    s_prev=None,
    mid_tail_weight_para=None,
):
    """
    第一阶段目标: log residual平方和 + 权重和的软约束 λ * (sum(w)-1)^2
    """
    ## 所有参数 >=0
    if np.any(params < 0):
        return 1e20
    ##
    s = params[n_gauss : 2 * n_gauss]
    if np.any(s <= 0):
        return 1e20

    y_pred = multi_gauss_model(x, params, n_gauss, with_background)
    eps = 1e-12
    if mid_tail_weight_para is None:
        W = 1.0
    else:
        W = mid_tail_weights(x, **mid_tail_weight_para)
    # weights = np.maximum(y, eps)
    # loss = np.sum(weights * (np.log(y_pred+eps) - np.log(y+eps))**2)
    loss = np.sum(W * np.pow(np.log(y_pred + eps) - np.log(y + eps), 2.0))
    # 软约束: 权重和接近1
    #    w_sum = np.sum(params[1:n_gauss+1])
    #    loss += lam_w * (w_sum - 1.0)**2
    # 软约束: normFactor接近1; 要求输入是归一化密度分布
    normFactor = params[0]
    loss += lam_n * np.pow(normFactor - 1.0, 2.0)
    ##
    if s_prev is not None:
        loss += sigma_smooth_penalty(s, s_prev)
    return loss


def objective_linear(
    params,
    x,
    y,
    n_gauss,
    with_background,
    lam_n=1e3,
    s_prev=None,
    mid_tail_weight_para=None,
):
    """
    第二阶段目标: linear residual平方和 + 权重和软约束 λ * (sum(w)-1)^2
    """
    if np.any(params < 0):
        return 1e20
    s = params[n_gauss : 2 * n_gauss]
    if np.any(s <= 0):
        return 1e20
    y_pred = multi_gauss_model(x, params, n_gauss, with_background)
    if mid_tail_weight_para is None:
        W = 1.0
    else:
        W = mid_tail_weights(x, **mid_tail_weight_para)
    loss = W * np.sum(np.pow(y_pred - y, 2.0))
    # 软约束: 权重和接近1
    #    w_sum = np.sum(params[1:n_gauss+1])
    #    loss += lam_w * (w_sum - 1.0)**2
    # 软约束: normFactor接近1
    normFactor = params[0]
    loss += lam_n * (normFactor - 1.0) ** 2
    if s_prev is not None:
        loss += sigma_smooth_penalty(s, s_prev)
    return loss


## scipy constraint; 等式约束表示约束函数结果为零，而不等式约束表示其为非负数
def build_constraints(n_gauss, is_post_bragg=False):
    cons = []
    # 不等式约束: sum(weight)<1, sigma递增
    # for i in range(n_gauss-1):
    #    cons.append({"type": "ineq", "fun": lambda p, i=i: p[i+1] - p[i+2]})
    ##================================== 归一约束
    cons.append({"type": "ineq", "fun": lambda p, n=n_gauss: 1 - np.sum(p[1:n])})
    ##================================== sigma 越来越大; 弥散约束
    for i in range(n_gauss - 1):
        if is_post_bragg and i == 0:
            continue
        cons.append({"type": "ineq", "fun": lambda p, i=i, n=n_gauss: p[n + i + 1] - p[n + i]})
    return cons


def build_bounds(x, n_gauss, s_prev=None, step_range=0.2, is_post_bragg=False, step_range_relax=5):
    ## + is append for list; * is duplicate
    lower = [0.5] + [1e-8] * (n_gauss - 1) + [1e-3] * n_gauss
    upper = [1.5] + [1.0] * (n_gauss - 1) + [(x.max() - x.min())] * n_gauss

    # 对于相邻层的拟合, 限制参数变化范围
    if s_prev is not None:
        for j in range(n_gauss):
            idx_sgm = 1 + (n_gauss - 1) + j  ## index of sigma
            # 判断是否是第一高斯sigma, 且在Bragg峰后
            # if j == 0 and is_post_bragg:
            if is_post_bragg:
                s_range = step_range_relax
            else:
                s_range = step_range

            lb = max(1e-3, (1.0 - s_range) * s_prev[j])
            ub = max(lb * 1.05, (1.0 + s_range) * s_prev[j])
            lower[idx_sgm] = lb
            upper[idx_sgm] = ub
    print(list(zip(lower, upper)))
    return list(zip(lower, upper))


def fit_multi_gauss(
    x,
    y,
    n_gauss=2,
    init_params=None,
    with_background=False,
    is_post_bragg=False,
    s_prev=None,
    lam_n=1e3,
    info_verbose=False,
    info_tag=None,
):
    x = np.asarray(x, float)
    y = np.asarray(y, float)

    # 初值
    if init_params is None:
        norm_init = 1.0
        w_init = np.ones(n_gauss - 1) * 0.5 / (n_gauss - 1)  # 平均分配权重
        # TODO weighted_stddev 在不是全范围时可能有问题?
        s_guess = weighted_stddev(x, y)
        s_init = [s_guess * (0.8 + 2 * i) for i in range(n_gauss)]
        # s_init = np.logspace(np.log10(1e-2), np.log10(30), n_gauss)
        p0 = np.array([norm_init] + list(w_init) + list(s_init) + ([y.min()] if with_background else []))
        print("初始param: ", p0)
    else:
        p0 = init_params.copy()

    #    # 边界
    #    lower = [0.5] + [1e-8]*(n_gauss-1) + [1e-3]*n_gauss
    #    upper = [1.5] + [1.0]*(n_gauss-1) + [5*(x.max()-x.min())]*n_gauss
    #    if with_background:
    #        lower.append(0)
    #        upper.append(0.5)
    #    bounds = list(zip(lower, upper))
    bounds = build_bounds(x, n_gauss, s_prev, is_post_bragg=is_post_bragg)

    constraints = build_constraints(n_gauss, is_post_bragg)

    ## 预拟合, 用来检测主峰两侧的误差, 以设置 mid_tail_weight
    res0 = minimize(
        objective_log,
        p0,
        args=(x, y, n_gauss, with_background, lam_n, s_prev, None),
        method="SLSQP",
        bounds=bounds,
        constraints=constraints,
        options={"maxiter": 500, "ftol": 1e-10},
    )
    y_pred0 = multi_gauss_model(x, res0.x, n_gauss, with_background)

    mid_tail_weight_para = mid_tail_detect_v2(x, y, y_pred0)
    ## print("mid_tail_weight_para: ", mid_tail_weight_para)
    if info_verbose:
        prefix = f"[{info_tag}] " if info_tag is not None else ""
        print(f"{prefix}mid_tail_weight_para:  center={mid_tail_weight_para['center']:.3f}, width={mid_tail_weight_para['width']:.3f}, lam_m={mid_tail_weight_para['lam_m']:.1f}, n={n_gauss}")

    # if n_gauss == 3:
    #    mid_tail_weight_para["lam_m"] = float(np.clip(1.5 * mid_tail_weight_para["lam_m"], 6.0, 40.0))

    ## 第一阶段: log residual拟合, 更关注尾部, 避免被主峰主导
    res1 = minimize(
        objective_log,
        p0,
        args=(x, y, n_gauss, with_background, lam_n),
        method="SLSQP",
        bounds=bounds,
        constraints=constraints,
        options={"maxiter": 3000, "ftol": 1e-10},
    )

    #    idx_tail_range = np.where(np.abs(x) > 3*res1.x[n_gauss+1])[0]

    #    res1 = minimize(objective_log, res1.x,
    #                      args=(x[idx_tail_range], y[idx_tail_range], n_gauss, with_background, lam_w, lam_n),
    #                      method="SLSQP", bounds=bounds, constraints=constraints,
    #                      options={"maxiter": 2000, "ftol": 1e-10})

    # 第二阶段: linear residual精修, 调整整体幅度
    res2 = minimize(
        objective_linear,
        res1.x,
        args=(x, y, n_gauss, with_background, lam_n),
        method="SLSQP",
        bounds=bounds,
        constraints=constraints,
        options={"maxiter": 2000, "ftol": 1e-10},
    )

    if not res2.success:
        print("优化未收敛: ", res2.message)

    return res2.x, res2


################################
# 开始拟合
#
################################


def main():
    # 创建输出目录
    enedir = "275MeV"
    output_dir = f"fit_plots_{enedir}"
    os.makedirs(f"{output_dir}", exist_ok=True)
    fit_results_txt = f"fit_results_{enedir}.txt"

    # prepare input data
    Zbin = 1000
    Zmax = 50.0
    Ybin = 1000
    Ymax = 50
    Ymin = -50
    Ymax_fit = 50
    Ymin_fit = -50

    Zaxis = np.arange(Zmax / Zbin / 2.0, Zmax, Zmax / Zbin)
    # print('Zaxis: {}'.format(Zaxis))
    Yaxis_raw = np.arange(Ymin + (Ymax - Ymin) / Ybin / 2, Ymax, (Ymax - Ymin) / Ybin)

    #     _2            _1              _3
    # BLOOD_ICRP, BONE_COMPACT_ICRU, BRAIN_ICRP
    # material = 'BLOOD_ICRP'
    # material = 'BONE_COMPACT_ICRU'
    # material = 'BRAIN_ICRP'
    material = "WATER"
    prtcls_name = "he"
    # enedir = '125MeV'
    dir_tag = "filter"

    # read .bin file and restore array shape
    path = "/data/wenxiao/cmcvalidate_letd_new/cudaCMC_validate_Geant4/G4Project/" + prtcls_name + "_" + enedir + "/"

    arr_he = np.fromfile(path + "eDep.bin", dtype=np.float64).reshape(Zbin, Ybin)

    IDD = arr_he.sum(axis=1)

    # N-高斯拟合设置
    n_gauss = 5
    results = np.zeros((Zbin, 2 * n_gauss + 1))  # 存储所有深度的拟合结果

    # 剂量最大位置
    idz0 = np.argmax(arr_he.sum(axis=1))
    peak_idd = IDD[idz0]

    # find peak first?
    # z_depth_index = int(idz0)
    z_depth_index = int(0.85 * idz0)  # 选一个peak之前的位置
    # z_depth_index = 78

    print("Zaxis[{}]: {}".format(z_depth_index, Zaxis[z_depth_index]))

    # 归一化数据
    profile = arr_he[z_depth_index, :]
    profile = normalize_by_sum(profile, (Ymax - Ymin) / Ybin)

    # ===================================
    # 限制拟合数据的横坐标范围
    mask = (Yaxis_raw >= Ymin_fit) & (Yaxis_raw <= Ymax_fit)
    Yaxis = Yaxis_raw[mask]
    profile = profile[mask]

    ## ===================================
    ## 先拟合布拉格峰位附近
    x, y = Yaxis, profile
    params, res = fit_multi_gauss(x, y, n_gauss, info_verbose=True, info_tag=f"idx={z_depth_index}")
    print("params", params)
    results[z_depth_index] = np.hstack(
        (
            [params[0]],
            [1 - np.sum(params[1:n_gauss])],
            params[1:n_gauss],
            params[n_gauss : 2 * n_gauss],
        )
    )
    print(
        "normFactor:",
        params[0],
        "第一高斯权重",
        1 - np.sum(params[1:n_gauss]),
        "其余拟合权重:",
        params[1:n_gauss],
        "sigma:",
        params[n_gauss : 2 * n_gauss],
        "其余权重和:",
        np.sum(params[1:n_gauss]),
    )
    y_fit = multi_gauss_model(x, params, n_gauss, False)

    y_min = np.min(profile)
    y_max = np.max(profile)
    y_lower = max(y_min * 0.5, 1e-12)
    y_upper = y_max * 2
    plt.figure(dpi=150)
    plt.semilogy(x, profile, "o", label="data", markersize=4, markerfacecolor="none", alpha=0.8)
    plt.semilogy(x, y_fit, "-", label="fit")
    plt.legend()
    plt.xlabel("x [mm]")
    plt.ylabel("Dose [a.u.]")
    plt.title(f"Depth idx={z_depth_index} (z={Zaxis[z_depth_index]:.3f} cm)")
    plt.ylim(y_lower, y_upper)
    plt.savefig(f"{output_dir}/depth_{z_depth_index:03d}.png")
    # plt.savefig(f"depth_{z_depth_index:03d}.png")
    plt.close()

    # 向更浅深度递推拟合
    for i in range(z_depth_index - 1, -1, -1):
        profile = normalize_by_sum(arr_he[i, :], (Ymax - Ymin) / Ybin)
        profile = profile[mask]
        s_prev = params[n_gauss : 2 * n_gauss]
        params, _ = fit_multi_gauss(
            x,
            profile,
            n_gauss,
            init_params=params,
            s_prev=s_prev,
            info_verbose=True,
            info_tag=f"idx={i}",
        )
        print("params", params)
        results[i] = np.hstack(
            (
                [params[0]],
                [1 - np.sum(params[1:n_gauss])],
                params[1:n_gauss],
                params[n_gauss : 2 * n_gauss],
            )
        )
        print(
            "normFactor:",
            params[0],
            "第一高斯权重",
            1 - np.sum(params[1:n_gauss]),
            "其余拟合权重:",
            params[1:n_gauss],
            "sigma:",
            params[n_gauss : 2 * n_gauss],
            "其余权重和:",
            np.sum(params[1:n_gauss]),
        )
        y_fit = multi_gauss_model(x, params, n_gauss, False)

        y_min = np.min(profile)
        y_max = np.max(profile)
        y_lower = max(y_min * 0.5, 1e-12)
        y_upper = y_max * 2
        # 绘图并保存
        plt.figure(dpi=150)
        plt.semilogy(x, profile, "o", label="data", markersize=4, markerfacecolor="none", alpha=0.8)
        plt.semilogy(x, y_fit, "-", label="fit")
        plt.legend()
        plt.xlabel("x [mm]")
        plt.ylabel("Dose [a.u.]")
        plt.title(f"Depth idx={i} (z={Zaxis[i]:.3f} cm)")
        plt.ylim(y_lower, y_upper)
        plt.savefig(f"{output_dir}/depth_{i:03d}.png")
        plt.close()

    # 重置params
    # params = results[z_depth_index]
    params = np.delete(results[z_depth_index], 1)

    # 向更深深度递推拟合, 有深度的截止条件
    # for i in range(z_depth_index+1, min(Zbin, 2*z_depth_index, z_depth_index+40)):
    for i in range(z_depth_index + 1, min(Zbin, idz0 + 40)):
        if IDD[i] < 0.01 * peak_idd:
            print(f"深度 {i} 已低于峰值 1%, 停止拟合. ")
            break

        profile = normalize_by_sum(arr_he[i, :], (Ymax - Ymin) / Ybin)
        profile = profile[mask]
        s_prev = params[n_gauss : 2 * n_gauss]
        if i > int(idz0):
            params, _ = fit_multi_gauss(
                Yaxis,
                profile,
                n_gauss,
                init_params=params,
                is_post_bragg=True,
                s_prev=s_prev,
                info_verbose=True,
                info_tag=f"idx={i}",
            )
        else:
            params, _ = fit_multi_gauss(
                Yaxis,
                profile,
                n_gauss,
                init_params=params,
                is_post_bragg=False,
                s_prev=s_prev,
                info_verbose=True,
                info_tag=f"idx={i}",
            )
        print("params", params)
        results[i] = np.hstack(
            (
                [params[0]],
                [1 - np.sum(params[1:n_gauss])],
                params[1:n_gauss],
                params[n_gauss : 2 * n_gauss],
            )
        )
        print(
            "normFactor:",
            params[0],
            "第一高斯权重",
            1 - np.sum(params[1:n_gauss]),
            "其余拟合权重:",
            params[1:n_gauss],
            "sigma:",
            params[n_gauss : 2 * n_gauss],
            "其余权重和:",
            np.sum(params[1:n_gauss]),
        )
        y_fit = multi_gauss_model(x, params, n_gauss, False)

        y_min = np.min(profile)
        y_max = np.max(profile)
        y_lower = max(y_min * 0.5, 1e-12)
        y_upper = y_max * 2
        # 绘图并保存
        plt.figure(dpi=150)
        plt.semilogy(x, profile, "o", label="data", markersize=4, markerfacecolor="none", alpha=0.8)
        plt.semilogy(x, y_fit, "-", label="fit")
        plt.legend()
        plt.xlabel("x [mm]")
        plt.ylabel("Dose [a.u.]")
        plt.title(f"Depth idx={i} (z={Zaxis[i]:.3f} cm)")
        plt.ylim(y_lower, y_upper)
        plt.savefig(f"{output_dir}/depth_{i:03d}.png")
        plt.close()

    # 保存结果到txt
    np.savetxt(f"{fit_results_txt}", results, header="norm, w1..wn, sigma1..n", fmt="%.6e")
    print(f"拟合完成, 结果已保存到 {fit_results_txt} 和 {output_dir}/*.png")


if __name__ == "__main__":
    main()
